﻿// CTreeListDlg.cpp: 实现文件
//

#include "pch.h"
#include "MY_PE_TOOL.h"
#include "CTreeListDlg.h"
#include "afxdialogex.h"


// CTreeListDlg 对话框

IMPLEMENT_DYNAMIC(CTreeListDlg, CDialogEx)

CTreeListDlg::CTreeListDlg(CWnd* pParent /*=nullptr*/)
	: CDialogEx(IDD_DIALOG6, pParent)
{

}

CTreeListDlg::~CTreeListDlg()
{
}

void CTreeListDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialogEx::DoDataExchange(pDX);
	DDX_Control(pDX, IDC_TREE1, m_TreeCtrl);
	DDX_Control(pDX, IDC_LIST1, m_ListCtrl);
}


BEGIN_MESSAGE_MAP(CTreeListDlg, CDialogEx)
	ON_WM_MOVE()
	ON_WM_SIZE()
    ON_NOTIFY(TVN_SELCHANGED, IDC_TREE1, &CTreeListDlg::OnSelchangedTree1)
END_MESSAGE_MAP()


// CTreeListDlg 消息处理程序


DWORD CTreeListDlg::RVATOFA(DWORD dwRVA)
{
    PIMAGE_DOS_HEADER pDosHead = (PIMAGE_DOS_HEADER)(*m_pFileData);
    PIMAGE_NT_HEADERS pNtHead = (PIMAGE_NT_HEADERS)((DWORD)pDosHead + pDosHead->e_lfanew);
    PIMAGE_OPTIONAL_HEADER32 pOptionalHead = &pNtHead->OptionalHeader;
    PIMAGE_SECTION_HEADER pSectionHead = (PIMAGE_SECTION_HEADER)((DWORD)&pNtHead->OptionalHeader + pNtHead->FileHeader.SizeOfOptionalHeader);

    //在节表中
    for (int i = 0; i < pNtHead->FileHeader.NumberOfSections; i++)
    {
        //判断该RVA是否在这个节区中
        if (dwRVA >= pSectionHead[i].VirtualAddress && dwRVA < pSectionHead[i].VirtualAddress + pSectionHead[i].SizeOfRawData)
        {
            return (dwRVA - pSectionHead[i].VirtualAddress + pSectionHead[i].PointerToRawData);
        }
    }

    return 0;
}


//展开树控件
void CTreeListDlg::ExpandTreeCtrl(HTREEITEM hTreeItem)
{
    HTREEITEM hNextItem = m_TreeCtrl.GetChildItem(hTreeItem);
    m_TreeCtrl.Expand(hTreeItem, TVE_EXPAND);

    while (hNextItem != NULL)
    {
        hNextItem = m_TreeCtrl.GetNextItem(hNextItem, TVGN_NEXT);    
        m_TreeCtrl.Expand(hTreeItem, TVE_EXPAND);
    }

}


void CTreeListDlg::ClearListCtrl()
{
    m_ListCtrl.DeleteAllItems();
    int nColumnCount = m_ListCtrl.GetHeaderCtrl()->GetItemCount();
    for (int i = 0; i < nColumnCount; i++)
    {
        m_ListCtrl.DeleteColumn(0);
    }
}

void CTreeListDlg::SetListCtrlColumn()
{
    m_ListCtrl.InsertColumn(0, _T("成员"), LVCFMT_LEFT, 150);
    m_ListCtrl.InsertColumn(1, _T("偏移量"), LVCFMT_LEFT, 150);
    m_ListCtrl.InsertColumn(2, _T("大小"), LVCFMT_LEFT, 150);
    m_ListCtrl.InsertColumn(3, _T("值"), LVCFMT_LEFT, 150);
}

CString CTreeListDlg::IntToString(DWORD nData, DWORD nLen)
{
    nLen *= 2;

    CString strFmt;
    strFmt.Format(_T("%%0%dX"), nLen);

    CString strRet;
    strRet.Format(strFmt, nData);

    return strRet;
}

CString CTreeListDlg::GetVariableSizeToString(DWORD nLen)
{
    switch (nLen)
    {
        case sizeof(BYTE) :
        {
            return _T("BYTE");
        }
        case sizeof(WORD) :
        {
            return _T("WORD");
        }
        case sizeof(DWORD) :
        {
            return _T("DWORD");
        }
    }

    return NULL;
}

void CTreeListDlg::ShowDirectory(_RESOURCE_INFO* pResInfo)
{
    int nIndex = 0;
    PIMAGE_RESOURCE_DIRECTORY pDirectory = (PIMAGE_RESOURCE_DIRECTORY)pResInfo->m_pAddr;

    m_ListCtrl.InsertItem(nIndex, _T("Characteristics"));
    m_ListCtrl.SetItemText(nIndex, 1, IntToString((DWORD) & (pDirectory->Characteristics) - (DWORD)(*m_pFileData), 4));
    m_ListCtrl.SetItemText(nIndex, 2, GetVariableSizeToString(sizeof(pDirectory->Characteristics)));
    m_ListCtrl.SetItemText(nIndex++, 3, IntToString(pDirectory->Characteristics, sizeof(pDirectory->Characteristics)));

    m_ListCtrl.InsertItem(nIndex, _T("TimeDateStamp"));
    m_ListCtrl.SetItemText(nIndex, 1, IntToString((DWORD) & (pDirectory->TimeDateStamp) - (DWORD)(*m_pFileData), 4));
    m_ListCtrl.SetItemText(nIndex, 2, GetVariableSizeToString(sizeof(pDirectory->TimeDateStamp)));
    m_ListCtrl.SetItemText(nIndex++, 3, IntToString(pDirectory->TimeDateStamp, sizeof(pDirectory->TimeDateStamp)));

    m_ListCtrl.InsertItem(nIndex, _T("MajorVersion"));
    m_ListCtrl.SetItemText(nIndex, 1, IntToString((DWORD) & (pDirectory->MajorVersion) - (DWORD)(*m_pFileData), 4));
    m_ListCtrl.SetItemText(nIndex, 2, GetVariableSizeToString(sizeof(pDirectory->MajorVersion)));
    m_ListCtrl.SetItemText(nIndex++, 3, IntToString(pDirectory->MajorVersion, sizeof(pDirectory->MajorVersion)));

    m_ListCtrl.InsertItem(nIndex, _T("MinorVersion"));
    m_ListCtrl.SetItemText(nIndex, 1, IntToString((DWORD) & (pDirectory->MinorVersion) - (DWORD)(*m_pFileData), 4));
    m_ListCtrl.SetItemText(nIndex, 2, GetVariableSizeToString(sizeof(pDirectory->MinorVersion)));
    m_ListCtrl.SetItemText(nIndex++, 3, IntToString(pDirectory->MinorVersion, sizeof(pDirectory->MinorVersion)));

    m_ListCtrl.InsertItem(nIndex, _T("NumberOfNamedEntries"));
    m_ListCtrl.SetItemText(nIndex, 1, IntToString((DWORD) & (pDirectory->NumberOfNamedEntries) - (DWORD)(*m_pFileData), 4));
    m_ListCtrl.SetItemText(nIndex, 2, GetVariableSizeToString(sizeof(pDirectory->NumberOfNamedEntries)));
    m_ListCtrl.SetItemText(nIndex++, 3, IntToString(pDirectory->NumberOfNamedEntries, sizeof(pDirectory->NumberOfNamedEntries)));

    m_ListCtrl.InsertItem(nIndex, _T("NumberOfIdEntries"));
    m_ListCtrl.SetItemText(nIndex, 1, IntToString((DWORD) & (pDirectory->NumberOfIdEntries) - (DWORD)(*m_pFileData), 4));
    m_ListCtrl.SetItemText(nIndex, 2, GetVariableSizeToString(sizeof(pDirectory->NumberOfIdEntries)));
    m_ListCtrl.SetItemText(nIndex++, 3, IntToString(pDirectory->NumberOfIdEntries, sizeof(pDirectory->NumberOfIdEntries)));
}


void CTreeListDlg::ShowDirectoryEntry(_RESOURCE_INFO* pResInfo)
{
    int nIndex = 0;
    PIMAGE_RESOURCE_DIRECTORY_ENTRY pDirectoryEntry = (PIMAGE_RESOURCE_DIRECTORY_ENTRY)pResInfo->m_pAddr;

    m_ListCtrl.InsertItem(nIndex, _T("Characteristics"));
    m_ListCtrl.SetItemText(nIndex, 1, IntToString((DWORD) & (pDirectoryEntry->Name) - (DWORD)(*m_pFileData), 4));
    m_ListCtrl.SetItemText(nIndex, 2, GetVariableSizeToString(sizeof(pDirectoryEntry->Name)));
    m_ListCtrl.SetItemText(nIndex++, 3, IntToString(pDirectoryEntry->Name, sizeof(pDirectoryEntry->Name)));

    m_ListCtrl.InsertItem(nIndex, _T("OffsetToData"));
    m_ListCtrl.SetItemText(nIndex, 1, IntToString((DWORD) & (pDirectoryEntry->OffsetToData) - (DWORD)(*m_pFileData), 4));
    m_ListCtrl.SetItemText(nIndex, 2, GetVariableSizeToString(sizeof(pDirectoryEntry->OffsetToData)));
    m_ListCtrl.SetItemText(nIndex++, 3, IntToString(pDirectoryEntry->OffsetToData, sizeof(pDirectoryEntry->OffsetToData)));
}

void CTreeListDlg::ShowDataEntry(_RESOURCE_INFO* pResInfo)
{
    int nIndex = 0;
    PIMAGE_RESOURCE_DATA_ENTRY pDataEntry = (PIMAGE_RESOURCE_DATA_ENTRY)pResInfo->m_pAddr;

    m_ListCtrl.InsertItem(nIndex, _T("OffsetToData"));
    m_ListCtrl.SetItemText(nIndex, 1, IntToString((DWORD) & (pDataEntry->OffsetToData) - (DWORD)(*m_pFileData), 4));
    m_ListCtrl.SetItemText(nIndex, 2, GetVariableSizeToString(sizeof(pDataEntry->OffsetToData)));
    m_ListCtrl.SetItemText(nIndex++, 3, IntToString(pDataEntry->OffsetToData, sizeof(pDataEntry->OffsetToData)));

    m_ListCtrl.InsertItem(nIndex, _T("Size"));
    m_ListCtrl.SetItemText(nIndex, 1, IntToString((DWORD) & (pDataEntry->Size) - (DWORD)(*m_pFileData), 4));
    m_ListCtrl.SetItemText(nIndex, 2, GetVariableSizeToString(sizeof(pDataEntry->Size)));
    m_ListCtrl.SetItemText(nIndex++, 3, IntToString(pDataEntry->Size, sizeof(pDataEntry->Size)));

    m_ListCtrl.InsertItem(nIndex, _T("CodePage"));
    m_ListCtrl.SetItemText(nIndex, 1, IntToString((DWORD) & (pDataEntry->CodePage) - (DWORD)(*m_pFileData), 4));
    m_ListCtrl.SetItemText(nIndex, 2, GetVariableSizeToString(sizeof(pDataEntry->CodePage)));
    m_ListCtrl.SetItemText(nIndex++, 3, IntToString(pDataEntry->CodePage, sizeof(pDataEntry->CodePage)));

    m_ListCtrl.InsertItem(nIndex, _T("Reserved"));
    m_ListCtrl.SetItemText(nIndex, 1, IntToString((DWORD) & (pDataEntry->Reserved) - (DWORD)(*m_pFileData), 4));
    m_ListCtrl.SetItemText(nIndex, 2, GetVariableSizeToString(sizeof(pDataEntry->Reserved)));
    m_ListCtrl.SetItemText(nIndex++, 3, IntToString(pDataEntry->Reserved, sizeof(pDataEntry->Reserved)));
}

void CTreeListDlg::SetFileData(char** pFileData)
{
    m_pFileData = pFileData;
}


void CTreeListDlg::ShowResourceData()
{
    m_TreeCtrl.DeleteAllItems();

    PIMAGE_DOS_HEADER pDosHead = (PIMAGE_DOS_HEADER)(*m_pFileData);
    PIMAGE_NT_HEADERS pNtHead = (PIMAGE_NT_HEADERS)((DWORD)pDosHead + pDosHead->e_lfanew);
    PIMAGE_OPTIONAL_HEADER32 pOptionalHead = &pNtHead->OptionalHeader;
    
    DWORD dwResourceFileOffset = RVATOFA(pOptionalHead->DataDirectory[IMAGE_DIRECTORY_ENTRY_RESOURCE].VirtualAddress);
    PIMAGE_RESOURCE_DIRECTORY pResource = (PIMAGE_RESOURCE_DIRECTORY)(dwResourceFileOffset + (DWORD)pDosHead);

    HTREEITEM hRoot = m_TreeCtrl.InsertItem(_T("Resource Directory"), TVI_ROOT);
    m_TreeCtrl.SetItemData(hRoot, (DWORD)new _RESOURCE_INFO((void*)pResource, REC_DIRECTORY));

    PIMAGE_RESOURCE_DIRECTORY_ENTRY pDirEntry = (PIMAGE_RESOURCE_DIRECTORY_ENTRY)((DWORD)pResource + sizeof(IMAGE_RESOURCE_DIRECTORY));
    ShowAllData(pDirEntry, hRoot, pResource->NumberOfNamedEntries + pResource->NumberOfIdEntries, (DWORD)pResource);

    //展开树控件
    ExpandTreeCtrl(hRoot);
}



void CTreeListDlg::ShowAllData(PIMAGE_RESOURCE_DIRECTORY_ENTRY pDirEntry, HTREEITEM hRoot, int nItem, DWORD SecOffset)
{
    for(int i = 0; i < nItem; i++)
    {
        CString strFmt;
        strFmt.Format(_T("Resource Directory Entry %d, ID: %d"), i + 1, pDirEntry[i].Id);
        HTREEITEM hTreeItem = m_TreeCtrl.InsertItem(strFmt, hRoot);
        m_TreeCtrl.SetItemData(hTreeItem, (DWORD)new _RESOURCE_INFO((void*)(pDirEntry + i), REC_DIRECTORY_ENTRY));

        if (pDirEntry[i].DataIsDirectory == 1)
        {
            PIMAGE_RESOURCE_DIRECTORY pResource = (PIMAGE_RESOURCE_DIRECTORY)(pDirEntry[i].OffsetToDirectory + SecOffset);
            HTREEITEM hTreeSubItem = m_TreeCtrl.InsertItem(_T("Resource Directory"), hTreeItem);
            m_TreeCtrl.SetItemData(hTreeSubItem, (DWORD)new _RESOURCE_INFO((void*)pResource, REC_DIRECTORY));

            PIMAGE_RESOURCE_DIRECTORY_ENTRY pNewDirEntry = (PIMAGE_RESOURCE_DIRECTORY_ENTRY)((DWORD)pResource + sizeof(IMAGE_RESOURCE_DIRECTORY));
            ShowAllData(pNewDirEntry, hTreeSubItem, pResource->NumberOfNamedEntries + pResource->NumberOfIdEntries, SecOffset);

        }
        else
        {
            HTREEITEM hTreeSubItem = m_TreeCtrl.InsertItem(_T("Resource Data Entry"), hTreeItem);
            m_TreeCtrl.SetItemData(hTreeSubItem, (DWORD)new _RESOURCE_INFO((void*)(SecOffset + pDirEntry[i].OffsetToDirectory), REC_DATA_ENTRY));
        }
    }
}



void CTreeListDlg::OnMove(int x, int y)
{
	CDialogEx::OnMove(x, y);

	// TODO: 在此处添加消息处理程序代码
}


void CTreeListDlg::OnSize(UINT nType, int cx, int cy)
{
	CDialogEx::OnSize(nType, cx, cy);

    if (m_TreeCtrl)
    {
        CRect rect;
        GetClientRect(rect);
        rect.bottom /= 2;
        rect.bottom -= 2;
        m_TreeCtrl.MoveWindow(rect);
    }
    if (m_ListCtrl)
    {
        CRect rect;
        GetClientRect(rect);
        rect.top = rect.bottom / 2 + 2;
        m_ListCtrl.MoveWindow(rect);
    }
}


BOOL CTreeListDlg::OnInitDialog()
{
    CDialogEx::OnInitDialog();



    return TRUE;  // return TRUE unless you set the focus to a control
                  // 异常: OCX 属性页应返回 FALSE
}


void CTreeListDlg::OnSelchangedTree1(NMHDR* pNMHDR, LRESULT* pResult)
{
    LPNMTREEVIEW pNMTreeView = reinterpret_cast<LPNMTREEVIEW>(pNMHDR);

    ClearListCtrl();
    SetListCtrlColumn();

    _RESOURCE_INFO* pResInfo = (_RESOURCE_INFO *)m_TreeCtrl.GetItemData(pNMTreeView->itemNew.hItem);

    switch (pResInfo->m_rType)
    {
    case REC_DIRECTORY:
    {
        ShowDirectory(pResInfo);
        break;
    }
    case REC_DIRECTORY_ENTRY:
    {
        ShowDirectoryEntry(pResInfo);
        break;
    }
    case REC_DATA_ENTRY:
    {
        ShowDataEntry(pResInfo);
        break;
    }
    }



    *pResult = 0;
}
